from keras.layers import Reshape, Conv2D, MaxPooling2D, Conv2DTranspose,
                         Flatten
from keras.datasets import mnist
from keras.models import Sequential

(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train = X_train.reshape(X_train.shape[0],
                          X_train.shape[1],X_train.shape[2],1)
X_test = X_test.reshape(X_test.shape[0],X_test.shape[1],X_test.shape[2],1)
X_train = X_train/255
X_test = X_test/255

model = Sequential()
model.add(Conv2D(16, (3,3), input_shape=(28, 28,1),
          activation='relu',padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2),name='pool1'))
model.add(Conv2D(8, (3,3), activation='relu',padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2),name='pool2'))
model.add(Conv2D(4, (3,3), activation='relu',padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2),name='pool3'))
model.add(Flatten(name='flatten'))
model.add(Reshape((3,3,4)))
model.add(Conv2DTranspose(4, kernel_size = (5,5), activation='relu'))
model.add(Conv2DTranspose(8, kernel_size = (8,8), activation='relu'))
model.add(Conv2DTranspose(16, kernel_size = (15,15), activation='relu'))
model.add(Conv2D(1, (3, 3), activation='relu',padding='same'))
model.summary()

model.compile(loss='mean_squared_error', optimizer='adam')
model.fit(X_train, X_train, validation_data=(X_test, X_test),epochs=10,
          batch_size=256, verbose=1)
X = []
for i in range(12):
    X.append(model.predict(X_test[i,:].reshape(1,28,28,1)).reshape(28,28))
plot(X)
